import gradio as gr
import os
import json
import uvicorn
from fastapi import FastAPI, HTTPException, UploadFile, File
from pydantic import BaseModel
from typing import Optional, List
from openai import OpenAI
import base64
import cv2
import numpy as np
from datetime import datetime

# --- 配置部分 ---
BAIDU_API_KEY = os.environ.get(
    "BAIDU_API_KEY", "xxx"  # tip: 请自行申请百度AI开放平台账号，并获取API Key
)
BAIDU_BASE_URL = "https://aistudio.baidu.com/llm/lmapi/v3"
MODEL_NAME = "ernie-4.5-turbo-vl"  # 百度多模态模型

# 初始化百度AI客户端（兼容OpenAI格式）
client = OpenAI(
    api_key=BAIDU_API_KEY,
    base_url=BAIDU_BASE_URL,
)


# --- 工具函数 ---
def encode_image_with_cv2(image_path: str) -> str:
    """使用OpenCV将图片编码为Base64字符串（支持本地文件路径）"""
    try:
        # 读取图片（处理中文路径问题）
        img = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        if img is None:
            raise ValueError(f"无法读取图片：{image_path}")

        # 编码为JPEG格式
        _, buffer = cv2.imencode(".jpg", img, [int(cv2.IMWRITE_JPEG_QUALITY), 85])
        return base64.b64encode(buffer).decode("utf-8")
    except Exception as e:
        print(f"图片编码错误：{str(e)}")
        raise


def encode_uploaded_image(file_content: bytes) -> str:
    """将上传的图片字节流编码为Base64字符串（用于API上传）"""
    try:
        img = cv2.imdecode(
            np.frombuffer(file_content, dtype=np.uint8), cv2.IMREAD_COLOR
        )
        _, buffer = cv2.imencode(".jpg", img, [int(cv2.IMWRITE_JPEG_QUALITY), 85])
        return base64.b64encode(buffer).decode("utf-8")
    except Exception as e:
        print(f"上传图片编码错误：{str(e)}")
        raise


# --- 核心业务函数 ---
def analyze_water_quality_core(
    parameters: dict,
    location: str,
    date: str,
    description: str = "",
    images: List[str] = None,
) -> dict:
    """水质分析核心逻辑（支持文本+图片输入）"""
    try:
        # 构建参数文本描述
        params_text = "\n".join([f"- {k}: {v}" for k, v in parameters.items()])
        prompt = f"""
你是资深水环境监测专家，需基于以下信息完成水质分析：
1. 基础信息：位置={location}，日期={date}，补充说明={description}
2. 监测参数：
{params_text}

请输出结构化分析报告，包含：
- 水质等级：明确对应《地表水环境质量标准》(GB 3838-2002) 中的类别
- 指标分析：标注异常指标及超标/达标情况
- 原因推断：基于参数异常提出可能的污染来源或影响因素
- 建议措施：针对性的治理、监测或防护建议
- 风险提示：对生态或人类健康的潜在影响（如有）
"""

        # 构建百度API要求的消息格式（文本+图片）
        messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
        if images and len(images) > 0:
            for img_base64 in images[:3]:  # 最多处理3张图片
                messages[0]["content"].append(
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{img_base64}"},
                    }
                )

        # 调用百度AI
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            stream=False,
            temperature=0.3,  # 低温度确保结果严谨
        )
        return {"success": True, "result": completion.choices[0].message.content}
    except Exception as e:
        error_msg = f"水质分析失败：{str(e)}"
        print(error_msg)
        return {"success": False, "error": error_msg}


def facility_diagnosis_core(
    facility_type: str,
    symptoms: str,
    usage_years: int,
    maintenance_history: str = "",
    images: List[str] = None,
) -> dict:
    """设施故障诊断核心逻辑（支持图片辅助诊断）"""
    try:
        prompt = f"""
你是公共设施维修专家，专注于{facility_type}的故障诊断。基于以下信息分析：
1. 设施信息：类型={facility_type}，使用年限={usage_years}年，维护记录={maintenance_history}
2. 故障症状：{symptoms}

请输出结构化诊断报告，包含：
- 故障定位：明确可能的故障部件或系统
- 原因分析：基于症状和使用年限推断故障根源
- 维修优先级：标注紧急程度（高/中/低）及影响范围
- 维修方案：分步骤的操作指南（含工具/材料建议）
- 预防建议：避免类似故障的日常维护措施
"""

        # 构建百度API消息
        messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
        if images and len(images) > 0:
            for img_base64 in images[:3]:
                messages[0]["content"].append(
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{img_base64}"},
                    }
                )

        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            stream=False,
            temperature=0.4,
        )
        return {"success": True, "result": completion.choices[0].message.content}
    except Exception as e:
        error_msg = f"设施诊断失败：{str(e)}"
        print(error_msg)
        return {"success": False, "error": error_msg}


def environmental_risk_assessment_core(
    project_type: str, location: str, description: str, sensitive_areas: str = ""
) -> dict:
    """环境风险评估核心逻辑（合规性导向）"""
    try:
        prompt = f"""
你是环境风险评估工程师，需基于以下项目信息完成风险分析：
1. 项目信息：类型={project_type}，地点={location}，项目描述={description}
2. 敏感区域：{sensitive_areas if sensitive_areas else '无'}

请输出符合《环境影响评价技术导则》的报告，包含：
- 风险识别：水/气/土/生态等维度的潜在风险项
- 风险等级：每项风险的可能性（高/中/低）及影响程度
- 合规性检查：是否符合相关环保法规（如《环境保护法》）
- mitigation措施：风险控制或缓解方案
- 监测计划：建议的环境监测指标和频率
"""

        # 构建百度API消息（纯文本）
        messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            stream=False,
            temperature=0.3,
        )
        return {"success": True, "result": completion.choices[0].message.content}
    except Exception as e:
        error_msg = f"风险评估失败：{str(e)}"
        print(error_msg)
        return {"success": False, "error": error_msg}


def generate_solution_core(
    problem: str, constraints: str, expected_outcome: str, related_data: str = ""
) -> dict:
    """解决方案生成核心逻辑（可行性导向）"""
    try:
        prompt = f"""
你是水利环境领域解决方案专家，需针对以下问题设计方案：
1. 问题描述：{problem}
2. 约束条件：{constraints if constraints else '无特殊约束'}
3. 期望结果：{expected_outcome}
4. 相关数据：{related_data if related_data else '无'}

请输出可落地的方案，包含：
- 方案框架：核心思路和技术路径
- 实施步骤：分阶段的执行计划（含时间节点）
- 资源需求：人力、设备、预算范围（估算）
- 效果验证：如何衡量方案是否达成目标
- 风险预案：实施中可能遇到的问题及应对策略
"""

        messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            stream=False,
            temperature=0.6,  # 稍高温度鼓励创新
        )
        return {"success": True, "result": completion.choices[0].message.content}
    except Exception as e:
        error_msg = f"方案生成失败：{str(e)}"
        print(error_msg)
        return {"success": False, "error": error_msg}


# --- Gradio界面包装函数 ---
def analyze_water_quality_for_gradio(
    ph, do, cod, bod, tn, tp, location, date, description, water_images
):
    """水质分析Gradio包装函数"""
    # 构建参数字典
    parameters = {
        "pH值": ph,
        "溶解氧(DO)": f"{do} mg/L",
        "化学需氧量(COD)": f"{cod} mg/L",
        "生化需氧量(BOD)": f"{bod} mg/L",
        "总氮(TN)": f"{tn} mg/L",
        "总磷(TP)": f"{tp} mg/L",
    }

    # 处理上传的图片
    encoded_images = []
    if water_images:
        for img_path in water_images:
            try:
                encoded = encode_image_with_cv2(img_path)
                encoded_images.append(encoded)
            except Exception as e:
                encoded_images.append(f"图片处理失败：{str(e)}")

    # 调用核心函数
    result = analyze_water_quality_core(
        parameters=parameters,
        location=location,
        date=date,
        description=description,
        images=encoded_images,
    )
    return result["result"] if result["success"] else result["error"]


def facility_diagnosis_for_gradio(
    facility_type, symptoms, usage_years, maintenance_history, facility_images
):
    """设施诊断Gradio包装函数"""
    encoded_images = []
    if facility_images:
        for img_path in facility_images:
            try:
                encoded = encode_image_with_cv2(img_path)
                encoded_images.append(encoded)
            except Exception as e:
                encoded_images.append(f"图片处理失败：{str(e)}")

    result = facility_diagnosis_core(
        facility_type=facility_type,
        symptoms=symptoms,
        usage_years=usage_years,
        maintenance_history=maintenance_history,
        images=encoded_images,
    )
    return result["result"] if result["success"] else result["error"]


def environmental_risk_for_gradio(project_type, location, description, sensitive_areas):
    """环境风险评估Gradio包装函数"""
    result = environmental_risk_assessment_core(
        project_type=project_type,
        location=location,
        description=description,
        sensitive_areas=sensitive_areas,
    )
    return result["result"] if result["success"] else result["error"]


def generate_solution_for_gradio(problem, constraints, expected_outcome, related_data):
    """解决方案生成Gradio包装函数"""
    result = generate_solution_core(
        problem=problem,
        constraints=constraints,
        expected_outcome=expected_outcome,
        related_data=related_data,
    )
    return result["result"] if result["success"] else result["error"]


# --- Gradio界面定义 ---
def gradio_interface():
    with gr.Blocks(
        title="水环智管师 - 水利环境AI助手",
        theme=gr.themes.Soft(
            primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.teal
        ),
    ) as demo:
        gr.Markdown("# 🌊 水环智管师 - 水利、环境与公共设施AI管理助手")
        gr.Markdown("基于文心一言4.5多模态能力，提供专业领域智能分析")

        with gr.Tabs():
            # 1. 水质分析标签页
            with gr.Tab("水质分析"):
                gr.Markdown("### 输入监测数据，获取GB标准合规分析")
                with gr.Row():
                    with gr.Column(scale=1):
                        ph = gr.Number(
                            label="pH值", value=7.0, precision=1, info="标准：6-9"
                        )
                        do = gr.Number(
                            label="溶解氧(DO)",
                            value=8.0,
                            precision=1,
                            info="单位：mg/L",
                        )
                        cod = gr.Number(
                            label="化学需氧量(COD)",
                            value=15.0,
                            precision=1,
                            info="单位：mg/L",
                        )
                    with gr.Column(scale=1):
                        bod = gr.Number(
                            label="生化需氧量(BOD)",
                            value=5.0,
                            precision=1,
                            info="单位：mg/L",
                        )
                        tn = gr.Number(
                            label="总氮(TN)", value=1.0, precision=2, info="单位：mg/L"
                        )
                        tp = gr.Number(
                            label="总磷(TP)", value=0.1, precision=2, info="单位：mg/L"
                        )
                with gr.Row():
                    location = gr.Textbox(
                        label="监测地点", placeholder="例：长江武汉段", scale=2
                    )
                    date = gr.Textbox(
                        label="监测日期",
                        value=datetime.now().strftime("%Y-%m-%d"),
                        scale=1,
                    )

                description = gr.Textbox(
                    label="补充说明", placeholder="例：雨后监测，有轻微浑浊"
                )
                water_images = gr.Files(
                    label="上传水样/监测点图片", file_types=["image"]
                )

                water_btn = gr.Button("生成水质分析报告", variant="primary")
                water_output = gr.Markdown(label="分析结果")
                water_btn.click(
                    fn=analyze_water_quality_for_gradio,
                    inputs=[
                        ph,
                        do,
                        cod,
                        bod,
                        tn,
                        tp,
                        location,
                        date,
                        description,
                        water_images,
                    ],
                    outputs=water_output,
                )

            # 2. 设施诊断标签页
            with gr.Tab("设施诊断"):
                gr.Markdown("### 描述故障症状，获取维修方案")
                facility_type = gr.Dropdown(
                    label="设施类型",
                    choices=[
                        "供水管道",
                        "污水处理设备",
                        "排水系统",
                        "水泵站",
                        "水质监测设备",
                        "防洪设施",
                        "其他",
                    ],
                    value="供水管道",
                    info="选择需要诊断的设施类别",
                )
                symptoms = gr.Textbox(
                    label="故障症状",
                    placeholder="例：管道接头处漏水，水压下降10%",
                    lines=2,
                )
                with gr.Row():
                    usage_years = gr.Number(
                        label="使用年限（年）", value=10, precision=0
                    )
                    maintenance_history = gr.Textbox(
                        label="维护历史", placeholder="例：2年前更换过密封圈"
                    )

                facility_images = gr.Files(
                    label="上传设施故障图片", file_types=["image"]
                )

                facility_btn = gr.Button("生成诊断报告", variant="primary")
                facility_output = gr.Markdown(label="诊断结果")
                facility_btn.click(
                    fn=facility_diagnosis_for_gradio,
                    inputs=[
                        facility_type,
                        symptoms,
                        usage_years,
                        maintenance_history,
                        facility_images,
                    ],
                    outputs=facility_output,
                )

            # 3. 环境风险评估标签页
            with gr.Tab("环境风险评估"):
                gr.Markdown("### 评估工程项目对环境的潜在影响")
                project_type = gr.Dropdown(
                    label="项目类型",
                    choices=[
                        "水利工程",
                        "污水处理厂建设",
                        "河道整治",
                        "市政管网改造",
                        "工业废水处理",
                        "城市供水系统",
                        "其他",
                    ],
                    value="水利工程",
                )
                location = gr.Textbox(
                    label="项目地点", placeholder="例：江苏省南京市长江沿岸"
                )
                description = gr.Textbox(
                    label="项目描述",
                    placeholder="例：新建一座日处理能力5万吨的污水处理厂，占地约20亩",
                    lines=2,
                )
                sensitive_areas = gr.Textbox(
                    label="周边敏感区域",
                    placeholder="例：下游3公里处有饮用水源保护区，西侧1公里为湿地公园",
                    lines=2,
                )

                risk_btn = gr.Button("生成风险评估报告", variant="primary")
                risk_output = gr.Markdown(label="评估结果")
                risk_btn.click(
                    fn=environmental_risk_for_gradio,
                    inputs=[project_type, location, description, sensitive_areas],
                    outputs=risk_output,
                )

            # 4. 解决方案生成标签页
            with gr.Tab("解决方案生成"):
                gr.Markdown("### 输入问题描述，获取专业解决方案")
                problem = gr.Textbox(
                    label="问题描述",
                    placeholder="例：某城市污水处理厂出水总磷浓度超标，无法达到排放标准",
                    lines=2,
                )
                constraints = gr.Textbox(
                    label="约束条件",
                    placeholder="例：改造预算不超过500万元，施工期间需保证70%处理能力",
                    lines=2,
                )
                expected_outcome = gr.Textbox(
                    label="期望结果",
                    placeholder="例：出水总磷稳定达到0.5mg/L以下，运行成本控制在合理范围",
                    lines=1,
                )
                related_data = gr.Textbox(
                    label="相关数据（可选）",
                    placeholder="例：当前处理工艺为A²/O，设计处理量10万吨/日，进水总磷平均3.2mg/L",
                    lines=2,
                )

                solution_btn = gr.Button("生成解决方案", variant="primary")
                solution_output = gr.Markdown(label="解决方案")
                solution_btn.click(
                    fn=generate_solution_for_gradio,
                    inputs=[problem, constraints, expected_outcome, related_data],
                    outputs=solution_output,
                )

        gr.Markdown(
            """
        ---
        ### 使用说明
        - 选择对应功能标签页，填写相关信息后点击按钮获取分析结果
        - 分析结果仅供参考，重要决策请咨询专业人员
        
        > 系统基于文心一言4.5多模态大模型开发，持续优化中
        """
        )

    return demo


# --- FastAPI应用和API端点定义 ---
app = FastAPI(
    title="水环智管师 - 水利环境AI服务",
    description="提供水利、环境和公共设施管理的AI分析服务",
)


# 请求模型定义
class WaterQualityRequest(BaseModel):
    parameters: dict
    location: str
    date: str
    description: Optional[str] = ""


class FacilityDiagnosisRequest(BaseModel):
    facility_type: str
    symptoms: str
    usage_years: int
    maintenance_history: Optional[str] = ""


class RiskAssessmentRequest(BaseModel):
    project_type: str
    location: str
    description: str
    sensitive_areas: Optional[str] = ""


class SolutionRequest(BaseModel):
    problem: str
    constraints: str
    expected_outcome: str
    related_data: Optional[str] = ""


# API端点
@app.post("/api/analyze-water-quality", summary="分析水质数据")
async def api_analyze_water_quality(request: WaterQualityRequest):
    result = analyze_water_quality_core(
        parameters=request.parameters,
        location=request.location,
        date=request.date,
        description=request.description,
    )
    if result["success"]:
        return result
    else:
        raise HTTPException(status_code=500, detail=result["error"])


@app.post("/api/diagnose-facility", summary="诊断设施故障")
async def api_diagnose_facility(request: FacilityDiagnosisRequest):
    result = facility_diagnosis_core(
        facility_type=request.facility_type,
        symptoms=request.symptoms,
        usage_years=request.usage_years,
        maintenance_history=request.maintenance_history,
    )
    if result["success"]:
        return result
    else:
        raise HTTPException(status_code=500, detail=result["error"])


@app.post("/api/assess-risk", summary="评估环境风险")
async def api_assess_risk(request: RiskAssessmentRequest):
    result = environmental_risk_assessment_core(
        project_type=request.project_type,
        location=request.location,
        description=request.description,
        sensitive_areas=request.sensitive_areas,
    )
    if result["success"]:
        return result
    else:
        raise HTTPException(status_code=500, detail=result["error"])


@app.post("/api/generate-solution", summary="生成解决方案")
async def api_generate_solution(request: SolutionRequest):
    result = generate_solution_core(
        problem=request.problem,
        constraints=request.constraints,
        expected_outcome=request.expected_outcome,
        related_data=request.related_data,
    )
    if result["success"]:
        return result
    else:
        raise HTTPException(status_code=500, detail=result["error"])


@app.post("/api/analyze-with-image", summary="带图片的综合分析")
async def api_analyze_with_image(
    analysis_type: str, text_data: str, files: List[UploadFile] = File(None)
):
    """通用的带图片分析接口"""
    try:
        # 编码图片
        encoded_images = []
        if files:
            for file in files:
                if "image" in file.content_type:
                    content = await file.read()
                    encoded = encode_uploaded_image(content)
                    encoded_images.append(encoded)

        # 解析文本数据
        data = json.loads(text_data)

        # 根据分析类型调用不同的核心函数
        if analysis_type == "water_quality":
            result = analyze_water_quality_core(
                parameters=data.get("parameters", {}),
                location=data.get("location", ""),
                date=data.get("date", ""),
                description=data.get("description", ""),
                images=encoded_images,
            )
        elif analysis_type == "facility_diagnosis":
            result = facility_diagnosis_core(
                facility_type=data.get("facility_type", ""),
                symptoms=data.get("symptoms", ""),
                usage_years=data.get("usage_years", 0),
                maintenance_history=data.get("maintenance_history", ""),
                images=encoded_images,
            )
        else:
            return {"success": False, "error": "不支持的分析类型"}

        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health", summary="健康检查")
async def health_check():
    return {
        "status": "healthy",
        "service": "水环智管师",
        "timestamp": datetime.now().isoformat(),
    }


# 挂载Gradio界面
demo = gradio_interface()
app = gr.mount_gradio_app(app=app, blocks=demo, path="/")

# 应用启动逻辑
if __name__ == "__main__":
    print("🚀 启动水环智管师服务...")
    print("📱 Gradio界面访问地址: http://localhost:7860")
    print("💻 FastAPI文档地址: http://localhost:7860/docs")
    print("-" * 50)

    uvicorn.run(app, host="0.0.0.0", port=7860)
